import numpy as np


# 数据预处理——数据归一化
class StandardScaler:
    def __init__(self):
        self.mean_ = None  # 均值
        self.scale_ = None  # 方差

    def fit(self, X):
        """根据训练数据集X获得数据的均值和标准差"""
        assert X.ndim == 2, "必须是二维数据"

        self.mean_ = np.array([np.mean(X[:, i]) for i in range(X.shape[1])])
        self.scale_ = np.array([np.std(X[:, i]) for i in range(X.shape[1])])

        return self

    def transform(self, X):
        """将X根据这个StandardScaler进行均值方差归一化处理"""
        assert X.ndim == 2, "必须是二维数据"
        assert self.mean_ is not None and self.scale_ is not None, "必须先调用fit方法传入训练数据"
        assert X.shape[1] == len(self.mean_), "待归一化的样本列数必须跟均值的列数相等"

        resX = np.empty(shape=X.shape, dtype=float)
        for col in range(X.shape[1]):
            resX[:, col] = (X[:, col] - self.mean_[col]) / self.scale_[col]

        return resX